# coding:utf-8
# @Author:Gscsd
from socketserver import ThreadingUDPServer, DatagramRequestHandler
from dnslib import DNSRecord, RR, A
import time
import logging

logging.basicConfig(
    filename="test.log",
    filemode="w",
    format="%(asctime)s %(name)s:%(levelname)s:%(message)s",
    datefmt="%d-%M-%Y %H:%M:%S",
    level=logging.DEBUG
)


class Hosts:
    max_age = 12 * 60 * 60
    record = {
        "drawin.bot": {"ip": '192.168.1.112',
                            "expires": 9e11},
        "drawinbot.com": {"ip": '192.168.1.112',
                            "expires": 9e11},
        "example.com": {"ip": '127.0.0.1',
                        "expires": 9e11}
    }
    cache = {}
    cache.update(record)

    @staticmethod
    def clear_cache():
        now_time = time.time()
        # 先进先出，一旦有未过期情况直接停止，节省资源
        for (k, v) in list(Hosts.cache.items()):
            if v.get("expires") <= now_time:
                Hosts.cache.pop(k)
            # 忽略自定义记录
            elif k in Hosts.record:
                continue
            else:
                break


class ThreadingUDPRequestHandler(DatagramRequestHandler):

    def handle(self):
        # 解析得到 DNSRecord 实例
        self.d = DNSRecord.parse(self.packet)
        # 获取查询类型
        self.qtype = self.d.q.qtype
        # qtype=1，即 A 记录
        if self.qtype == 1:
            # 构造响应
            r = self.d.reply()
            # 添加 DNS 解析回复
            r.add_answer(self.get_answer_A())
            # 封包，写入响应数据
            self.wfile.write(r.pack())
        else:
            # 向 114 DNS请求得到非 A 记录响应
            self.wfile.write(self.d.send('114.114.114.114', 53))

    def finish(self):
        super().finish()
        # A 记录则在响应客户端后检查缓存
        if self.qtype == 1:
            Hosts.clear_cache()
            logging.info(f'A -- {self.domain}, from {self.client_address}, to {self.ip}')

    def get_answer_A(self) -> RR:
        # TYPE=A，则 qname 是其对应的域名
        self.domain = str(self.d.q.qname).strip('.')
        now_time = time.time()
        if self.domain in Hosts.cache:
            self.ip = Hosts.cache.get(self.domain).get("ip")
            print(self.ip)
            return RR(self.domain, self.qtype, rdata=A(self.ip), ttl=60)
        r = DNSRecord.parse(self.d.send('114.114.114.114', 53))
        # TYPE=A，则 rdata 是其对应的 ip
        self.ip = r.a.rdata
        print(r.a)
        print(self.ip)
        print('-' * 20)
        # 添加记录到缓存
        # Hosts.cache.update({self.domain: {
        #     "ip": str(self.ip),
        #     "expires": now_time + Hosts.max_age
        # }})
        return r.a


if __name__ == "__main__":
    HOST, PORT = "0.0.0.0", 53
    with ThreadingUDPServer((HOST, PORT), ThreadingUDPRequestHandler) as server:
        # 持续运行
        server.serve_forever()